from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any, Literal, Union

# --- Base Model for common fields ---
class BaseNode(BaseModel):
    """Base model for all knowledge graph nodes."""
    full_path: str = Field(..., description="Fully qualified path (e.g., my_project.module.Class.method)")
    docstring: Optional[str] = Field(None, description="Docstring of the element")
    # Pass 2 Fields - Initialized empty, filled by the linker
    referenced_by: List[str] = Field(default_factory=list, description="List of full_paths that reference this node")

# --- Specific Node Types ---

class ModuleNode(BaseNode):
    """Represents a Python module (.py file)."""
    type: Literal['module'] = 'module'
    # file_path: str = Field(..., description="Absolute file system path to the .py file")
    # Pass 1 Fields
    imports: List[str] = Field(default_factory=list, description="List of full_paths imported by this module")
    functions: List[str] = Field(default_factory=list, description="List of full_paths of functions defined in this module")
    classes: List[str] = Field(default_factory=list, description="List of full_paths of classes defined in this module")
    variables: List[str] = Field(default_factory=list, description="List of full_paths of global variables defined in this module")
    # Optional: module_body_calls / module_body_references if needed for top-level code execution

class FunctionNode(BaseNode):
    """Represents a standalone function."""
    type: Literal['function'] = 'function'
    name: str = Field(..., description="Local name of the function")
    module_full_path: str = Field(..., description="Full path of the module containing this function")
    # Pass 1 Fields
    signature: Optional[str] = Field(None, description="Function signature (e.g., '(x: int, y: str) -> bool')")
    return_type_hint: Optional[str] = Field(None, description="Return type hint as a string")
    implementation: Optional[str] = Field(None, description="Source code of the function body")
    decorators: List[str] = Field(default_factory=list, description="List of full_paths of decorators applied")
    calls: List[str] = Field(default_factory=list, description="List of full_paths of functions/methods called by this function")
    references: List[str] = Field(default_factory=list, description="List of full_paths of variables/classes referenced")
    nested_classes: List[str] = Field(default_factory=list, description="List of full_paths of classes defined inside this function")
    nested_functions: List[str] = Field(default_factory=list, description="List of full_paths of functions defined inside this function")
    # Pass 2 Fields
    called_by: List[str] = Field(default_factory=list, description="List of full_paths of functions/methods that call this function")
    # Optional: raises: List[str] = Field(default_factory=list)

class MethodNode(BaseNode):
    """Represents a method within a class."""
    type: Literal['method'] = 'method'
    name: str = Field(..., description="Local name of the method")
    class_name: str = Field(..., description="Local name of the owning class")
    class_full_path: str = Field(..., description="Full path of the owning class")
    method_type: Literal['instance', 'class', 'static'] = Field(..., description="Type of the method")
    # Pass 1 Fields
    signature: Optional[str] = Field(None, description="Method signature")
    return_type_hint: Optional[str] = Field(None, description="Return type hint as a string")
    implementation: Optional[str] = Field(None, description="Source code of the method body")
    decorators: List[str] = Field(default_factory=list, description="List of full_paths of decorators applied")
    calls: List[str] = Field(default_factory=list, description="List of full_paths of functions/methods called by this method")
    references: List[str] = Field(default_factory=list, description="List of full_paths of variables/classes/attributes referenced")
    nested_classes: List[str] = Field(default_factory=list, description="List of full_paths of classes defined inside this method")
    nested_functions: List[str] = Field(default_factory=list, description="List of full_paths of functions defined inside this method")
    # Pass 2 Fields
    called_by: List[str] = Field(default_factory=list, description="List of full_paths of functions/methods that call this method")
    # Optional: raises: List[str] = Field(default_factory=list)


# --- Supporting models for ClassNode ---

class ClassVariableInfo(BaseModel):
    """Information about a class variable."""
    name: str = Field(..., description="Name of the class variable")
    full_path: str = Field(..., description="Full path of the class variable (e.g., my_module.MyClass.VAR)")
    value_repr: Optional[str] = Field(None, description="String representation of the variable's value (best effort)")
    type_hint: Optional[str] = Field(None, description="Type hint annotation as a string")
    docstring: Optional[str] = Field(None, description="Docstring/comment associated with the variable")
    # Pass 2 Fields (Optional, could track references specifically to class vars)
    # referenced_by: List[str] = Field(default_factory=list)

class InstanceVariableInfo(BaseModel):
    """Information about an instance variable (inferred)."""
    name: str = Field(..., description="Name of the instance variable (attribute name)")
    full_path: str = Field(..., description="Conceptual full path (e.g., my_module.MyClass.attribute)") # Note: Not directly addressable usually
    type_hint: Optional[str] = Field(None, description="Type hint annotation if available (e.g., from __init__)")
    docstring: Optional[str] = Field(None, description="Docstring/comment if found near assignment")
    defined_in: List[str] = Field(default_factory=list, description="List of full_paths of methods where this variable is assigned (e.g., __init__)")
     # Pass 2 Fields (Optional, could track references specifically to instance vars)
    # referenced_by: List[str] = Field(default_factory=list)

# --- Class Node ---

class ClassNode(BaseNode):
    """Represents a class definition."""
    type: Literal['class'] = 'class'
    name: str = Field(..., description="Local name of the class")
    module_full_path: str = Field(..., description="Full path of the module containing this class")
    # Pass 1 Fields
    decorators: List[str] = Field(default_factory=list, description="List of full_paths of decorators applied to the class")
    inherits_from: List[str] = Field(default_factory=list, description="List of full_paths of direct base classes")
    methods: List[str] = Field(default_factory=list, description="List of full_paths of methods defined in this class")
    properties: List[str] = Field(default_factory=list, description="List of full_paths of properties")
    class_variables: List[ClassVariableInfo] = Field(default_factory=list, description="Details of class variables")
    instance_variables: List[InstanceVariableInfo] = Field(default_factory=list, description="Inferred details of instance variables")
    nested_classes: List[str] = Field(default_factory=list, description="List of full_paths of classes defined inside this class")
    nested_functions: List[str] = Field(default_factory=list, description="List of full_paths of functions defined inside this class (rare)")
    slots: Optional[List[str]] = Field(None, description="List of slot names if __slots__ is defined")
    references: List[str] = Field(default_factory=list, description="External symbols referenced directly in the class body (rare)")
    calls: List[str] = Field(default_factory=list, description="Function calls made directly in the class body (rare)")
    # Pass 2 Fields (referenced_by is inherited from BaseNode, covers instantiation, annotation etc.)
    # Optional: mro: List[str] = Field(default_factory=list) # Method Resolution Order (Pass 1 or 2)
    # Optional: metaclass: Optional[str] = None # (Pass 1)

# --- Variable Node ---

class VariableNode(BaseNode):
    """Represents a module-level (global) variable/constant."""
    type: Literal['variable'] = 'variable'
    name: str = Field(..., description="Name of the variable")
    module_full_path: str = Field(..., description="Full path of the module containing this variable")
    # Pass 1 Fields
    value_repr: Optional[str] = Field(None, description="String representation of the variable's value (best effort)")
    type_hint: Optional[str] = Field(None, description="Type hint annotation as a string")
    # is_constant: Optional[bool] = None # Based on heuristics (e.g., ALL_CAPS name)
    # Pass 2 Fields (referenced_by is inherited from BaseNode)
    # Optional: assigned_in: List[str] # Where assignment happens (usually just the module)


# --- Union Type for the final list ---
# This allows type checkers to know the list can contain any of these node types.
CodeKnowledgeNode = Union[ModuleNode, FunctionNode, MethodNode, ClassNode, VariableNode]


# Example Usage (can be removed later or moved to tests)
if __name__ == '__main__':
    print("Pydantic models for Code Knowledge Graph nodes defined.")

    # Example instantiation (won't have real data yet)
    module_ex = ModuleNode(
        full_path="my_project.data_utils",
        file_path="/abs/path/to/my_project/data_utils.py",
        docstring="Module docstring.",
        imports=["pandas", "my_project.config"],
        functions=["my_project.data_utils.load_data"],
        classes=[],
        variables=["my_project.data_utils.DEFAULT_ENCODING"]
    )

    function_ex = FunctionNode(
        type='function', # Redundant here as Pydantic handles it, but shows structure
        name='load_data',
        full_path='my_project.data_utils.load_data',
        module_full_path='my_project.data_utils',
        signature='(file_path: str) -> pandas.DataFrame',
        return_type_hint='pandas.DataFrame',
        docstring='Loads data.',
        calls=['pandas.read_csv'],
        references=['my_project.config.DEFAULT_CHUNK_SIZE']
    )

    print("\nExample ModuleNode:")
    print(module_ex.model_dump_json(indent=2)) # Use model_dump_json in Pydantic v2

    print("\nExample FunctionNode:")
    print(function_ex.model_dump_json(indent=2))

    # Example showing default empty lists for Pass 2 fields
    print(f"\nFunction called_by (initially): {function_ex.called_by}")
    print(f"Module referenced_by (initially): {module_ex.referenced_by}")